import torch
import torch.nn as nn
from torchvision.models import (
    resnet50, ResNet50_Weights,
    resnet18, ResNet18_Weights,
    resnet101, ResNet101_Weights,
    vgg16, VGG16_Weights,
    densenet121, DenseNet121_Weights,
    mobilenet_v2, MobileNet_V2_Weights
)

from PublicModel import MyModel


class ImageResNet50(MyModel):
    def __init__(self, num_classes=1000, pretrained=False, device=None):
        super().__init__()
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")

        # 使用 torchvision 的 ResNet50 权重管理方式
        weights = ResNet50_Weights.DEFAULT if pretrained else None
        self.model = resnet50(weights=weights)

        # 替换最后一层全连接层
        in_features = self.model.fc.in_features
        self.model.fc = nn.Linear(in_features, num_classes)

        self.model = self.model.to(self.device)

    def forward(self, x):
        # x: [batch, 3, H, W]
        x = x.to(self.device)
        return self.model(x)


class ImageResNet50(MyModel):
    def __init__(self, num_classes=1000, pretrained=False, device=None):
        super().__init__()
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        weights = ResNet50_Weights.DEFAULT if pretrained else None
        self.model = resnet50(weights=weights)
        self.model.fc = nn.Linear(self.model.fc.in_features, num_classes)
        self.model = self.model.to(self.device)

    def forward(self, x):
        return self.model(x.to(self.device))

class ImageResNet18(MyModel):
    def __init__(self, num_classes=1000, pretrained=False, device=None):
        super().__init__()
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        weights = ResNet18_Weights.DEFAULT if pretrained else None
        self.model = resnet18(weights=weights)
        self.model.fc = nn.Linear(self.model.fc.in_features, num_classes)
        self.model = self.model.to(self.device)

    def forward(self, x):
        return self.model(x.to(self.device))

class ImageResNet101(MyModel):
    def __init__(self, num_classes=1000, pretrained=False, device=None):
        super().__init__()
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        weights = ResNet101_Weights.DEFAULT if pretrained else None
        self.model = resnet101(weights=weights)
        self.model.fc = nn.Linear(self.model.fc.in_features, num_classes)
        self.model = self.model.to(self.device)

    def forward(self, x):
        return self.model(x.to(self.device))

class ImageVGG16(MyModel):
    def __init__(self, num_classes=1000, pretrained=False, device=None):
        super().__init__()
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        weights = VGG16_Weights.DEFAULT if pretrained else None
        self.model = vgg16(weights=weights)
        self.model.classifier[6] = nn.Linear(4096, num_classes)
        self.model = self.model.to(self.device)

    def forward(self, x):
        return self.model(x.to(self.device))

class ImageDenseNet121(MyModel):
    def __init__(self, num_classes=1000, pretrained=False, device=None):
        super().__init__()
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        weights = DenseNet121_Weights.DEFAULT if pretrained else None
        self.model = densenet121(weights=weights)
        self.model.classifier = nn.Linear(self.model.classifier.in_features, num_classes)
        self.model = self.model.to(self.device)

    def forward(self, x):
        return self.model(x.to(self.device))

class ImageMobileNetV2(MyModel):
    def __init__(self, num_classes=1000, pretrained=False, device=None):
        super().__init__()
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        weights = MobileNet_V2_Weights.DEFAULT if pretrained else None
        self.model = mobilenet_v2(weights=weights)
        self.model.classifier[1] = nn.Linear(self.model.classifier[1].in_features, num_classes)
        self.model = self.model.to(self.device)

    def forward(self, x):
        return self.model(x.to(self.device))
